# coding=utf-8
'''
利用百度baidu-aip进行人脸对齐
'''
import os, cv2, numpy
import logging
import sys
import time
import numpy as np
import dlib
from baidu_analyze_img import get_landmark
import argparse

parser = argparse.ArgumentParser(description='其他参数')
parser.add_argument('indir', type=str, help='输入文件夹')
parser.add_argument('outdir', type=str, help='输出文件夹')
parser.add_argument('--imgSize', type=list, default=[256, 256], help='设置shape')
parser.add_argument('--points5', type=list, default=[[92, 116],
                                                     [178, 112],
                                                     [133, 147],
                                                     [100, 172],
                                                     [154, 154]], help='五点标志')
parser.add_argument('--gender', default=None, type=str, help="截取男女性别（男为male，女为female，不填则全部截取）")
args = parser.parse_args()
imgSize = args.imgSize
coord5point = args.points5
indir = args.indir
outdir = args.outdir

# log的层级建立
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    filename=os.path.join("log", "man_" + str(time.time()).replace(".", "_") + ".log")
)

'''获取照片的几个位置'''


def getrect(path):
    predictor = dlib.shape_predictor("dat/shape_predictor_68_face_landmarks.dat")
    face = cv2.imread(path)
    rec = dlib.rectangle(0, 0, face.shape[0], face.shape[1])
    shape = predictor(np.uint8(face), rec)
    # left eye, right eye, nose, left mouth, right mouth
    order = [36, 45, 30, 48, 54]
    rect = []
    for j in order:
        x = shape.part(j).x
        y = shape.part(j).y
        rect.append([x, y])
    return rect


def transformation_from_points(points1, points2):
    points1 = points1.astype(numpy.float64)
    points2 = points2.astype(numpy.float64)
    c1 = numpy.mean(points1, axis=0)
    c2 = numpy.mean(points2, axis=0)
    points1 -= c1
    points2 -= c2
    s1 = numpy.std(points1)
    s2 = numpy.std(points2)
    points1 /= s1
    points2 /= s2
    U, S, Vt = numpy.linalg.svd(points1.T * points2)
    R = (U * Vt).T
    return numpy.vstack([numpy.hstack(((s2 / s1) * R, c2.T - (s2 / s1) * R * c1.T)), numpy.matrix([0., 0., 1.])])


def warp_im(img_im, orgi_landmarks, tar_landmarks):
    pts1 = numpy.float64(numpy.matrix([[point[0], point[1]] for point in orgi_landmarks]))
    pts2 = numpy.float64(numpy.matrix([[point[0], point[1]] for point in tar_landmarks]))
    M = transformation_from_points(pts1, pts2)
    dst = cv2.warpAffine(img_im, M[:2], (img_im.shape[1], img_im.shape[0]))
    return dst


def main_fa2(img_in, imgSize, coord5point, face_landmarks, dirout):
    pic_path = img_in
    img_im = cv2.imread(pic_path)
    img_im = cv2.resize(img_im, tuple(imgSize))
    dst = warp_im(img_im, face_landmarks, coord5point)
    crop_im = dst[0:imgSize[0], 0:imgSize[1]]
    crop_im = cv2.resize(crop_im, tuple(imgSize))
    ss = str(time.time()).replace(".", "_") + ".jpg"
    ans_path = os.path.join(dirout, ss)
    cv2.imwrite(ans_path, crop_im)


def catch_face(imgin, dirout):
    '''
    对一张照片进行截取
    :param imgin: path of img
    :param dirout: out path of imgdir
    :return:
    '''


def catch_faces(dirin, dirout):
    '''
    对于文件夹内的所有照片进行截取
    :param dirin: path of in dir
    :param dirout: path of out dir
    :return:
    '''
    print("开始执行")
    if os.path.exists(dirin) is False:
        print("文件夹不存在")
        return
    path, name = os.path.split(dirout)
    if os.path.exists(os.path.join(path, name + "male")) is False:
        os.mkdir(os.path.join(path, name + "male"))
    if os.path.exists(os.path.join(path, name + "female")) is False:
        os.mkdir(os.path.join(path, name + "female"))
    for root, dirs, files in os.walk(dirin, topdown=False):
        for file in files:
            try:
                imgin = os.path.join(root, file)
                face_landmarks_male, face_landmarks_female = get_landmark(imgin)
                path, name = os.path.split(dirout)
                if len(face_landmarks_male):
                    for face_landmarks in face_landmarks_male:
                        main_fa2(imgin, imgSize, coord5point, face_landmarks, os.path.join(path, name + "male"))
                        print(f"成功转换{imgin}")
                if len(face_landmarks_female):
                    for face_landmarks in face_landmarks_female:
                        main_fa2(imgin, imgSize, coord5point, face_landmarks, os.path.join(path, name + "female"))
                        print(f"成功转换{imgin}")
            except Exception as e:
                logging.error(f"转换失败:{path}")
                print(e)
                print(f"转换失败:{path}")


if __name__ == '__main__':
    catch_faces(indir, outdir)
